/*
 * Copyright 2016 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#ifndef GrVkResourceProvider_DEFINED
#define GrVkResourceProvider_DEFINED

#include "include/gpu/vk/GrVkTypes.h"
#include "include/private/base/SkMutex.h"
#include "include/private/base/SkTArray.h"
#include "src/base/SkTInternalLList.h"
#include "src/core/SkChecksum.h"
#include "src/core/SkLRUCache.h"
#include "src/core/SkTDynamicHash.h"
#include "src/gpu/ganesh/GrGpu.h"
#include "src/gpu/ganesh/GrManagedResource.h"
#include "src/gpu/ganesh/GrProgramDesc.h"
#include "src/gpu/ganesh/GrResourceHandle.h"
#include "src/gpu/ganesh/GrThreadSafePipelineBuilder.h"
#include "src/gpu/ganesh/vk/GrVkDescriptorPool.h"
#include "src/gpu/ganesh/vk/GrVkDescriptorSetManager.h"
#include "src/gpu/ganesh/vk/GrVkPipelineStateBuilder.h"
#include "src/gpu/ganesh/vk/GrVkRenderPass.h"
#include "src/gpu/ganesh/vk/GrVkSampler.h"
#include "src/gpu/ganesh/vk/GrVkSamplerYcbcrConversion.h"
#include "src/gpu/ganesh/vk/GrVkUtil.h"

class GrVkCommandPool;
class GrVkGpu;
class GrVkPipeline;
class GrVkPipelineState;
class GrVkPrimaryCommandBuffer;
class GrVkRenderTarget;
class GrVkSecondaryCommandBuffer;
class GrVkUniformHandler;

class GrVkResourceProvider {
public:
    GrVkResourceProvider(GrVkGpu *gpu);
    ~GrVkResourceProvider();

    GrThreadSafePipelineBuilder *pipelineStateCache()
    {
        return fPipelineStateCache.get();
    }

    sk_sp<GrThreadSafePipelineBuilder> refPipelineStateCache()
    {
        return fPipelineStateCache;
    }

    // Set up any initial vk objects
    void init();

    sk_sp<const GrVkPipeline> makePipeline(const GrProgramInfo &, VkPipelineShaderStageCreateInfo *shaderStageInfo,
        int shaderStageCount, VkRenderPass compatibleRenderPass, VkPipelineLayout layout, uint32_t subpass);

    GR_DEFINE_RESOURCE_HANDLE_CLASS(CompatibleRPHandle)

    using SelfDependencyFlags = GrVkRenderPass::SelfDependencyFlags;
    using LoadFromResolve = GrVkRenderPass::LoadFromResolve;

    // Finds or creates a simple render pass that matches the target, increments the refcount,
    // and returns. The caller can optionally pass in a pointer to a CompatibleRPHandle. If this is
    // non null it will be set to a handle that can be used in the furutre to quickly return a
    // compatible GrVkRenderPasses without the need inspecting a GrVkRenderTarget.
    const GrVkRenderPass *findCompatibleRenderPass(GrVkRenderTarget *target, CompatibleRPHandle *compatibleHandle,
        bool withResolve, bool withStencil, SelfDependencyFlags selfDepFlags, LoadFromResolve);
    const GrVkRenderPass *findCompatibleRenderPass(GrVkRenderPass::AttachmentsDescriptor *,
        GrVkRenderPass::AttachmentFlags, SelfDependencyFlags selfDepFlags, LoadFromResolve,
        CompatibleRPHandle *compatibleHandle = nullptr);

    const GrVkRenderPass *findCompatibleExternalRenderPass(VkRenderPass, uint32_t colorAttachmentIndex);


    // Finds or creates a render pass that matches the target and LoadStoreOps, increments the
    // refcount, and returns. The caller can optionally pass in a pointer to a CompatibleRPHandle.
    // If this is non null it will be set to a handle that can be used in the future to quickly
    // return a GrVkRenderPass without the need to inspect a GrVkRenderTarget.
    // TODO: sk_sp?
    const GrVkRenderPass *findRenderPass(GrVkRenderTarget *target, const GrVkRenderPass::LoadStoreOps &colorOps,
        const GrVkRenderPass::LoadStoreOps &resolveOps, const GrVkRenderPass::LoadStoreOps &stencilOps,
        CompatibleRPHandle *compatibleHandle, bool withResolve, bool withStencil, SelfDependencyFlags selfDepFlags,
        LoadFromResolve);

    // The CompatibleRPHandle must be a valid handle previously set by a call to findRenderPass or
    // findCompatibleRenderPass.
    const GrVkRenderPass *findRenderPass(const CompatibleRPHandle &compatibleHandle,
        const GrVkRenderPass::LoadStoreOps &colorOps, const GrVkRenderPass::LoadStoreOps &resolveOps,
        const GrVkRenderPass::LoadStoreOps &stencilOps);

    GrVkCommandPool *findOrCreateCommandPool();

    void checkCommandBuffers();

    void forceSyncAllCommandBuffers();

    // We must add the finishedProc to all active command buffers since we may have flushed work
    // that the client cares about before they explicitly called flush and the GPU may reorder
    // command execution. So we make sure all previously submitted work finishes before we call the
    // finishedProc.
    void addFinishedProcToActiveCommandBuffers(sk_sp<skgpu::RefCntedCallback> finishedCallback);

    // Finds or creates a compatible GrVkDescriptorPool for the requested type and count.
    // The refcount is incremented and a pointer returned.
    // TODO: Currently this will just create a descriptor pool without holding onto a ref itself
    //       so we currently do not reuse them. Rquires knowing if another draw is currently using
    //       the GrVkDescriptorPool, the ability to reset pools, and the ability to purge pools out
    //       of our cache of GrVkDescriptorPools.
    GrVkDescriptorPool *findOrCreateCompatibleDescriptorPool(VkDescriptorType type, uint32_t count);

    // Finds or creates a compatible GrVkSampler based on the GrSamplerState and
    // GrVkYcbcrConversionInfo. The refcount is incremented and a pointer returned.
    GrVkSampler *findOrCreateCompatibleSampler(GrSamplerState, const GrVkYcbcrConversionInfo &ycbcrInfo);

    // Finds or creates a compatible GrVkSamplerYcbcrConversion based on the GrSamplerState and
    // GrVkYcbcrConversionInfo. The refcount is incremented and a pointer returned.
    GrVkSamplerYcbcrConversion *findOrCreateCompatibleSamplerYcbcrConversion(const GrVkYcbcrConversionInfo &ycbcrInfo);

    GrVkPipelineState *findOrCreateCompatiblePipelineState(GrRenderTarget *, const GrProgramInfo &,
        VkRenderPass compatibleRenderPass, bool overrideSubpassForResolveLoad);

    GrVkPipelineState *findOrCreateCompatiblePipelineState(const GrProgramDesc &, const GrProgramInfo &,
        VkRenderPass compatibleRenderPass, GrThreadSafePipelineBuilder::Stats::ProgramCacheResult *stat);

    sk_sp<const GrVkPipeline> findOrCreateMSAALoadPipeline(const GrVkRenderPass &renderPass, int numSamples,
        VkPipelineShaderStageCreateInfo *, VkPipelineLayout);

    void getSamplerDescriptorSetHandle(VkDescriptorType type, const GrVkUniformHandler &,
        GrVkDescriptorSetManager::Handle *handle);

    // This is a convenience function to return a descriptor set for zero sammples. When making a
    // VkPipelineLayout we must pass in an array of valid descriptor set handles. However, we have
    // set up our system to have the descriptor sets be in the order uniform, sampler, input. So
    // if we have a uniform and input we will need to have a valid handle for the sampler as well.
    // When using the GrVkMSAALoadManager this is the case, but we also don't have a
    // GrVkUniformHandler to pass into the more general function. Thus we use this call instead.
    void getZeroSamplerDescriptorSetHandle(GrVkDescriptorSetManager::Handle *handle);

    // Returns the compatible VkDescriptorSetLayout to use for uniform buffers. The caller does not
    // own the VkDescriptorSetLayout and thus should not delete it. This function should be used
    // when the caller needs the layout to create a VkPipelineLayout.
    VkDescriptorSetLayout getUniformDSLayout() const;

    // Returns the compatible VkDescriptorSetLayout to use for input attachments. The caller does
    // not own the VkDescriptorSetLayout and thus should not delete it. This function should be used
    // when the caller needs the layout to create a VkPipelineLayout.
    VkDescriptorSetLayout getInputDSLayout() const;

    // Returns the compatible VkDescriptorSetLayout to use for a specific sampler handle. The caller
    // does not own the VkDescriptorSetLayout and thus should not delete it. This function should be
    // used when the caller needs the layout to create a VkPipelineLayout.
    VkDescriptorSetLayout getSamplerDSLayout(const GrVkDescriptorSetManager::Handle &) const;

    // Returns a GrVkDescriptorSet that can be used for uniform buffers. The GrVkDescriptorSet
    // is already reffed for the caller.
    const GrVkDescriptorSet *getUniformDescriptorSet();

    // Returns a GrVkDescriptorSet that can be used for sampler descriptors that are compatible with
    // the GrVkDescriptorSetManager::Handle passed in. The GrVkDescriptorSet is already reffed for
    // the caller.
    const GrVkDescriptorSet *getSamplerDescriptorSet(const GrVkDescriptorSetManager::Handle &);

    // Returns a GrVkDescriptorSet that can be used for input attachments. The GrVkDescriptorSet
    // is already reffed for the caller.
    const GrVkDescriptorSet *getInputDescriptorSet();

    // Signals that the descriptor set passed it, which is compatible with the passed in handle,
    // can be reused by the next allocation request.
    void recycleDescriptorSet(const GrVkDescriptorSet *descSet, const GrVkDescriptorSetManager::Handle &);

    void storePipelineCacheData();

    // Destroy any cached resources. To be called before destroying the VkDevice.
    // The assumption is that all queues are idle and all command buffers are finished.
    // For resource tracing to work properly, this should be called after unrefing all other
    // resource usages.
    void destroyResources();

    // Currently we just release available command pools (which also releases their buffers). The
    // command buffers and pools take up the most memory. Other objects (e.g. samples,
    // ycbcr conversions, etc.) tend to be fairly light weight and not worth the effort to remove
    // them and then possibly remake them. Additionally many of those objects have refs/handles that
    // are held by other objects that aren't deleted here. Thus the memory wins for removing these
    // objects from the cache are probably not worth the complexity of safely releasing them.
    void releaseUnlockedBackendObjects();

#if defined(GR_TEST_UTILS)
    void resetShaderCacheForTesting() const
    {
        fPipelineStateCache->release();
    }
#endif

private:
    class PipelineStateCache : public GrThreadSafePipelineBuilder {
    public:
        PipelineStateCache(GrVkGpu *gpu);
        ~PipelineStateCache() override;

        void release();
        GrVkPipelineState *findOrCreatePipelineState(GrRenderTarget *, const GrProgramInfo &,
            VkRenderPass compatibleRenderPass, bool overrideSubpassForResolveLoad);
        GrVkPipelineState *findOrCreatePipelineState(const GrProgramDesc &desc, const GrProgramInfo &programInfo,
            VkRenderPass compatibleRenderPass, Stats::ProgramCacheResult *stat)
        {
            return this->findOrCreatePipelineStateImpl(desc, programInfo, compatibleRenderPass, false, stat);
        }

    private:
        struct Entry;

        GrVkPipelineState *findOrCreatePipelineStateImpl(const GrProgramDesc &, const GrProgramInfo &,
            VkRenderPass compatibleRenderPass, bool overrideSubpassForResolveLoad, Stats::ProgramCacheResult *);

        struct DescHash {
            uint32_t operator () (const GrProgramDesc &desc) const
            {
                return SkChecksum::Hash32(desc.asKey(), desc.keyLength());
            }
        };

        SkLRUCache<const GrProgramDesc, std::unique_ptr<Entry>, DescHash> fMap;

        GrVkGpu *fGpu;
    };

    class CompatibleRenderPassSet {
    public:
        // This will always construct the basic load store render pass (all attachments load and
        // store their data) so that there is at least one compatible VkRenderPass that can be used
        // with this set.
        CompatibleRenderPassSet(GrVkRenderPass *renderPass);

        bool isCompatible(const GrVkRenderPass::AttachmentsDescriptor &, GrVkRenderPass::AttachmentFlags,
            SelfDependencyFlags selfDepFlags, LoadFromResolve) const;

        const GrVkRenderPass *getCompatibleRenderPass() const
        {
            // The first GrVkRenderpass should always exist since we create the basic load store
            // render pass on create
            SkASSERT(fRenderPasses[0]);
            return fRenderPasses[0];
        }

        GrVkRenderPass *getRenderPass(GrVkGpu *gpu, const GrVkRenderPass::LoadStoreOps &colorOps,
            const GrVkRenderPass::LoadStoreOps &resolveOps, const GrVkRenderPass::LoadStoreOps &stencilOps);

        void releaseResources();

    private:
        skia_private::STArray<4, GrVkRenderPass *> fRenderPasses;
        int fLastReturnedIndex;
    };

    VkPipelineCache pipelineCache();

    GrVkGpu *fGpu;

    // Central cache for creating pipelines
    VkPipelineCache fPipelineCache;

    struct MSAALoadPipeline {
        sk_sp<const GrVkPipeline> fPipeline;
        const GrVkRenderPass *fRenderPass;
    };

    // Cache of previously created msaa load pipelines
    skia_private::TArray<MSAALoadPipeline> fMSAALoadPipelines;

    skia_private::STArray<4, CompatibleRenderPassSet> fRenderPassArray;

    skia_private::TArray<const GrVkRenderPass *> fExternalRenderPasses;

    // Array of command pools that we are waiting on
    skia_private::STArray<4, GrVkCommandPool *, true> fActiveCommandPools;

    // Array of available command pools that are not in flight
    skia_private::STArray<4, GrVkCommandPool *, true> fAvailableCommandPools;

    // Stores GrVkSampler objects that we've already created so we can reuse them across multiple
    // GrVkPipelineStates
    SkTDynamicHash<GrVkSampler, GrVkSampler::Key> fSamplers;

    // Stores GrVkSamplerYcbcrConversion objects that we've already created so we can reuse them.
    SkTDynamicHash<GrVkSamplerYcbcrConversion, GrVkSamplerYcbcrConversion::Key> fYcbcrConversions;

    // Cache of GrVkPipelineStates
    sk_sp<PipelineStateCache> fPipelineStateCache;

    skia_private::STArray<4, std::unique_ptr<GrVkDescriptorSetManager>> fDescriptorSetManagers;

    GrVkDescriptorSetManager::Handle fUniformDSHandle;
    GrVkDescriptorSetManager::Handle fInputDSHandle;
};

#endif
